import torch
from torch import nn
from torch.nn import functional as F

import numpy as np

from detectron2.config import configurable
from detectron2.layers import Conv2d, ShapeSpec, get_norm
from detectron2.structures import Boxes, ImageList, Instances, BitMasks

class ROIEnhance():
    def __init__(self,
                shift:bool=False,
                expand:bool=False
                ) -> None:
        self._shift = shift
        self._expand = expand
        self._scale = 4
        pass
    
    def getBoxesofTargetCFG(self,target):
        tgt_masks = target["masks"]
        tgt_ids = target["gt_ids"]
        boxes = BitMasks(tgt_masks).get_bounding_boxes()
        
        if self._shift:
            _,imw,imh = target['masks'].shape
            boxes_ts = self._rand_shift(boxes.tensor,imw,imh)
        
        if self._expand:
            _,imw,imh = target['masks'].shape
            bs_ts = self._expand_box(boxes.tensor,imw,imh)
            boxes.tensor = bs_ts
        
        valid_box_index = boxes.nonempty(threshold=8.0)
        valid_boxes = boxes[valid_box_index]
        valid_boxes = np.ceil(valid_boxes.tensor / self._scale ).to(torch.int)
        return boxes,valid_boxes,valid_box_index
        
    def getBoxesinTargetsCFG(self,targets):
        b_tensor = torch.tensor([],dtype=torch.int)
        ids_tensor = torch.tensor([],dtype=torch.int)
        boxes_tensor = torch.tensor([],dtype=torch.int)
        _,imw,imh = targets[0]['masks'].shape
        bs = len(targets)
        for b in range(bs):
            tgt_masks = targets[b]["masks"]
            tgt_ids = targets[b]["gt_ids"]
            boxes = BitMasks(tgt_masks).get_bounding_boxes()
            valid_box = boxes.nonempty(threshold=8.0)
            
            boxes = boxes[valid_box]
            tgt_ids = tgt_ids[valid_box].to(ids_tensor)

            if(len(boxes) < 1):
                continue
            
            boxes_ts = boxes.tensor
            if self._shift:
                boxes_ts = self._rand_shift(boxes.tensor,imw,iwh,scale=0.5)
            # if self._expand:
            #     boxes_ts = self._expand_box(boxes.tensor,imw,imh,scale=0.5)  
                
            boxes4 = np.ceil(boxes_ts / self._scale)
            boxes4= boxes4.to(torch.int)
            teb = torch.full((len(boxes),), b)
            
            boxes_tensor = torch.cat((boxes_tensor,boxes4),dim = 0)
            ids_tensor = torch.cat((ids_tensor,tgt_ids),dim=0)
            b_tensor = torch.cat((b_tensor,teb),dim=0)
        
        return boxes_tensor,b_tensor,ids_tensor
    
    def _rand_shift(self,boxes_tensor,maxw,maxh,scale=1.0):
        boxes_ts = boxes_tensor
        ww = boxes_ts[:,2] - boxes_ts[:,0]
        hh = boxes_ts[:,3] - boxes_ts[:,1]
        
        dw = torch.sqrt(ww)
        dh = torch.sqrt(hh)
        sn = torch.rand(size=(dw.shape[0],2)) * 2 - 1
        dw1 = dw * sn[:,0]
        dh1 = dh * sn[:,1]
        dwh = torch.zeros(size=(dw.shape[0],4))
        dwh[:,0] = dw1
        dwh[:,1] = dh1
        dwh[:,2] = dw1
        dwh[:,3] = dh1
        
        boxes_ts_2= boxes_ts + dwh * scale

        for inds in range(boxes_ts.shape[0]):
            if (boxes_ts_2[inds,0] < 0 or boxes_ts_2[inds,1] < 0 or boxes_ts_2[inds,2] > maxw or boxes_ts_2[inds,3] > maxh) :
                boxes_ts_2[inds] = boxes_ts[inds]
        return boxes_ts_2
    
    def _expand_box(self,boxes_ts,max_w,max_h,scale=1):
        ww = boxes_ts[:,2] - boxes_ts[:,0]
        hh = boxes_ts[:,3] - boxes_ts[:,1]
        dw = torch.sqrt(ww)*scale//2
        dh = torch.sqrt(hh)*scale//2
        
        ts0 = boxes_ts[:,0]
        ts1 = boxes_ts[:,1]
        ts2 = boxes_ts[:,2]
        ts3 = boxes_ts[:,3]
        ts0 = torch.where(ts0 > dw,ts0-dw,ts0)
        ts1 = torch.where(ts1 > dh,ts1-dh,ts1)
        ts2 = torch.where(ts2+dw < max_w,ts2+dw,ts2)
        ts3 = torch.where(ts3+dh < max_h,ts3+dh,ts3)
        
        boxes_ts[:,0] = ts0
        boxes_ts[:,1] = ts1
        boxes_ts[:,2] = ts2
        boxes_ts[:,3] = ts3
        return boxes_ts
    
    def getBoxesinTargetsWithShift(self,targets):
        b_tensor = torch.tensor([],dtype=torch.int)
        ids_tensor = torch.tensor([],dtype=torch.int)
        boxes_tensor = torch.tensor([],dtype=torch.int)
        _,imw,imh = targets[0]['masks'].shape
        bs = len(targets)
        for b in range(bs):
            tgt_masks = targets[b]["masks"]
            tgt_ids = targets[b]["gt_ids"]
            boxes = BitMasks(tgt_masks).get_bounding_boxes()
            
            valid_box = boxes.nonempty(threshold=8.0)
            
            boxes = boxes[valid_box]
            tgt_ids = tgt_ids[valid_box].to(ids_tensor)
            if(len(boxes) < 1):
                continue
            
            boxes_ts  = boxes.tensor
            ww = boxes_ts[:,2] - boxes_ts[:,0]
            hh = boxes_ts[:,3] - boxes_ts[:,1]
            
            dw = torch.sqrt(ww)
            dh = torch.sqrt(hh)
            sn = torch.rand(size=(dw.shape[0],2)) * 2 - 1
            dw1 = dw * sn[:,0]
            dh1 = dh * sn[:,1]
            dwh = torch.zeros(size=(dw.shape[0],4))
            dwh[:,0] = dw1
            dwh[:,1] = dh1
            dwh[:,2] = dw1
            dwh[:,3] = dh1
            
            boxes_ts_2= boxes_ts + dwh
            boxes_ts_3= boxes_ts - dwh

            for inds in range(boxes_ts.shape[0]):
                if (boxes_ts_2[inds,0] < 0 or boxes_ts_2[inds,1] < 0 or boxes_ts_2[inds,2] > imw or boxes_ts_2[inds,3] > imh) :
                    boxes_ts_2[inds] = boxes_ts[inds]
                if (boxes_ts_3[inds,0] < 0 or boxes_ts_3[inds,1] < 0 or boxes_ts_3[inds,2] > imw or boxes_ts_3[inds,3] > imh) :
                    boxes_ts_3[inds] = boxes_ts[inds]
                    
            boxes_ts_2 = torch.cat((boxes_ts_2,boxes_ts_3),dim = 0)
            boxes4 = np.ceil(boxes_ts_2 / 4 )
            boxes4= boxes4.to(torch.int)
            tgt_ids = torch.cat((tgt_ids,tgt_ids),dim=0)
            
            teb = torch.full((len(boxes) * 2,), b)
            boxes_tensor = torch.cat((boxes_tensor,boxes4),dim = 0)
            ids_tensor = torch.cat((ids_tensor,tgt_ids),dim=0)
            b_tensor = torch.cat((b_tensor,teb),dim=0)
        
        return boxes_tensor,b_tensor,ids_tensor
    
    def getBoxesofTarget(self,targets,bnum):
        tgt_masks = targets[bnum]["masks"]
        tgt_ids = targets[bnum]["gt_ids"]
        boxes = BitMasks(tgt_masks).get_bounding_boxes()
        valid_box = boxes.nonempty(threshold=8.0)
        boxes4 = np.ceil(boxes.tensor / 4 )
        boxes4= boxes4.to(torch.int)
        return boxes4,tgt_ids,valid_box

    
    def getBoxesinTargets(self,targets,scale = 8):
        b_tensor = torch.tensor([],dtype=torch.int)
        ids_tensor = torch.tensor([],dtype=torch.int)
        boxes_tensor = torch.tensor([],dtype=torch.int)
        
        bs = len(targets)
        for b in range(bs):
            tgt_masks = targets[b]["masks"]
            tgt_ids = targets[b]["gt_ids"]
            boxes = BitMasks(tgt_masks).get_bounding_boxes()
            valid_box = boxes.nonempty(threshold=8.0)
            
            boxes = boxes[valid_box]
            tgt_ids = tgt_ids[valid_box].to(ids_tensor)

            if(len(boxes) < 1):
                continue
            boxes4 = np.ceil(boxes.tensor / scale)
            boxes4= boxes4.to(torch.int)
            teb = torch.full((len(boxes),), b)
            
            boxes_tensor = torch.cat((boxes_tensor,boxes4),dim = 0)
            ids_tensor = torch.cat((ids_tensor,tgt_ids),dim=0)
            b_tensor = torch.cat((b_tensor,teb),dim=0)
        
        return boxes_tensor,b_tensor,ids_tensor

    def getBoxesinPred(self,preds):
        pass

# -------------------------------------------------------------------------------

def print_nb_params(m):
    model_parameters = filter(lambda p: p.requires_grad, m.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])
    print(f'Trainable parameters: {params/1e6:.3}M')


def main():
    x = torch.randn(1, 1024, 20, 20)
    # agg = MixVPR(
    #     in_channels=1024,
    #     in_h=20,
    #     in_w=20,
    #     out_channels=1024,
    #     mix_depth=4,
    #     mlp_ratio=1,
    #     out_rows=4)

    # print_nb_params(agg)
    # output = agg(x)
    # print(output.shape)


if __name__ == '__main__':
    main()
